# migrations/update_tables_0928.py
import os
import sys

# 添加项目根目录到Python路径
current_file = os.path.abspath(__file__)
stock_data_system_dir = os.path.dirname(os.path.dirname(current_file))
stock_traing_root = os.path.dirname(stock_data_system_dir)

if stock_traing_root not in sys.path:
    sys.path.insert(0, stock_traing_root)

from sqlalchemy import text
from stock_data_system.config.database import DatabaseConfig
from stock_data_system.database.postgres_manager import PostgresManager

def upgrade():
    """添加缺失的列到现有表"""

    # 创建数据库配置
    config = DatabaseConfig(
        host=os.getenv('DB_HOST', 'localhost'),
        port=os.getenv('DB_PORT', '5432'),
        database=os.getenv('DB_NAME', 'stock_data'),
        username=os.getenv('DB_USER', 'postgres'),
        password=os.getenv('DB_PASSWORD', 'password')
    )

    # 创建PostgresManager实例
    db_manager = PostgresManager(config)

    # 检查数据库连接
    if not db_manager.check_connection():
        print("数据库连接失败，请检查配置")
        return False

    try:
        with db_manager.engine.connect() as conn:
            print("开始更新数据库表结构...")

            # 为 stock_basic 表添加新列
            print("更新 stock_basic 表...")
            conn.execute(text("""
                ALTER TABLE stock_basic 
                ADD COLUMN IF NOT EXISTS type VARCHAR(20),
                ADD COLUMN IF NOT EXISTS country VARCHAR(10) DEFAULT 'CN',
                ADD COLUMN IF NOT EXISTS currency VARCHAR(10) DEFAULT 'CNY',
                ADD COLUMN IF NOT EXISTS delisting_date DATE
            """))

            # 为 daily_quotes 表添加新列
            print("更新 daily_quotes 表...")
            conn.execute(text("""
                ALTER TABLE daily_quotes 
                ADD COLUMN IF NOT EXISTS data_source VARCHAR(20) DEFAULT 'baostock'
            """))

            # 为 financial_data 表添加新列
            print("更新 financial_data 表...")
            conn.execute(text("""
                ALTER TABLE financial_data 
                ADD COLUMN IF NOT EXISTS total_profit FLOAT,
                ADD COLUMN IF NOT EXISTS operating_profit FLOAT,
                ADD COLUMN IF NOT EXISTS revenue_growth_rate FLOAT,
                ADD COLUMN IF NOT EXISTS profit_growth_rate FLOAT,
                ADD COLUMN IF NOT EXISTS asset_growth_rate FLOAT,
                ADD COLUMN IF NOT EXISTS equity_growth_rate FLOAT,
                ADD COLUMN IF NOT EXISTS eps_growth_rate FLOAT,
                ADD COLUMN IF NOT EXISTS asset_turnover FLOAT,
                ADD COLUMN IF NOT EXISTS inventory_turnover FLOAT,
                ADD COLUMN IF NOT EXISTS receivable_turnover FLOAT,
                ADD COLUMN IF NOT EXISTS current_assets FLOAT,
                ADD COLUMN IF NOT EXISTS current_liabilities FLOAT,
                ADD COLUMN IF NOT EXISTS asset_liability_ratio FLOAT,
                ADD COLUMN IF NOT EXISTS data_source VARCHAR(20) DEFAULT 'baostock'
            """))

            # 为 dividend_data 表添加新列
            print("更新 dividend_data 表...")
            conn.execute(text("""
                ALTER TABLE dividend_data 
                ADD COLUMN IF NOT EXISTS data_source VARCHAR(20) DEFAULT 'baostock'
            """))

            conn.commit()
            print("数据库表结构更新完成！")
            return True

    except Exception as e:
        print(f"更新数据库表结构失败: {e}")
        return False

if __name__ == "__main__":
    upgrade()